import pydicom as dicom
import h5py
import os
import nibabel
import nrrd
import numpy as np
from pydicom.dataset import Dataset, FileMetaDataset, FileDataset
import SimpleITK as sitk


# 读取功能相关的工具
class Load:
    def __init__(self, path, post):
        self.PATH = path
        self.POST = post
        self.lstFiles = []

    # 该函数是为了按照需要，从指定路径中读取所需的数据，为了降低内存的需求，所以存放的是路径
    def load_scan(self):
        patients = os.listdir(self.PATH)
        # 为了避免数据是无序的，故先按照文件名进行排序
        patients.sort()
        # 遍历指定路径下的文件，dirName即path，subdirList以及fileList分别是存放该路径下的子文件夹与文件名字的列表
        for dirName, subdirList, fileList in os.walk(self.PATH):
            for filename in fileList:
                # 将文件中后缀名为post指定的文件名加到路径中，然后存入lstFiles，并返回
                if self.POST in filename.lower():
                    self.lstFiles.append(os.path.join(dirName, filename))
        return self.lstFiles


# H5相关工具
class H5pyStore:
    def __init__(self, path, _type):
        self.PATH = path
        self.TYPE = _type
        self.FILE = h5py.File(self.PATH+'.hdf5', self.TYPE)

    def createDataSet(self, name, shape_or_array, _type=None):
        # 通过两种方式可以创建数据集
        if _type is not None:
            dataset = self.FILE.create_dataset(name, shape_or_array, _type)
        else:
            dataset = self.FILE.create_dataset(name, data=shape_or_array)
        # 打印输出FILE里的每个数据库
        print("/n已有数据集如下：")
        for key in self.FILE.keys():
            print(key, "的信息: Name：", self.FILE[key].name, "  Shape：", self.FILE[key].shape, "  Value：", self.FILE[key])
        return dataset

    def setDataValue(self, key, value):
        # 利用键值对赋值
        self.FILE[key] = value
        # 或者可以利用以下形式赋值
        # ds = self.FILE[key]
        # ds[...] = value
        return value

    def createGroup(self, name, parent=None):
        if parent is None:
            self.FILE.create_group(name)
        else:
            self.FILE[parent].creat_group(name)

    def close(self):
        self.FILE.close()

    def keys(self):
        return self.FILE.keys()

    def get(self, key):
        return self.FILE[key]


# 文件格式互转的工具
class Trans:
    def __init__(self):
        pass

    def trans(self, src, tar, name=None):
        if '.dcm' in src or 'dcm' == tar:
            self.transDCM(src, tar, name)
        else:
            self.transNoDCM(src, tar, name)

    def transNoDCM(self, src, tar, name=None):
        s = src[src.rindex('.')+1:]
        if s == 'gz':
            s = 'nii'
        t = tar if tar != 'gz' else 'nii'
        if name is None:
            name = src[:src.rindex('.')]
            if '.nii' in name:
                name = src[:src.rindex('.')]
        tar = name + '.' + t
        if s == t:
            print('无需转换格式')
        else:
            print(src, '--->', tar)
            self.__getattribute__(s+'TO'+t)(src, tar)
        return True

    def nrrdTOnii(self, src, tar):
        # 读出NRRD文件的内容
        data, header = nrrd.read(src)
        # np.eye(4)是生成一个4*4的单位矩阵，进行映射
        img = nibabel.Nifti1Image(data, np.eye(4))

        print(img.shape)
        nibabel.save(img, tar)

    def niiTOnrrd(self, src, tar):
        data = nibabel.load(src).get_fdata()
        nrrd.write(tar, data)

    def transDCM(self, src, tar, name=None):
        if tar == 'dcm':
            s = src[src.rindex('.')+1:]
            f = self.__getattribute__(s+'TOdcm')
        else:
            f = self.__getattribute__('dcmTO'+tar)
        if name is None:
            if tar == 'dcm':
                name = src[:src.rindex('.')] + '/'
            else:
                if '.dcm' in src:
                    name = src[:src.rindex('.')+1]+tar
                elif src[-1] == '/' or src[-1] == '\\':
                    name = src+'.'+tar
                else:
                    src += '/'
                    name = src+'.'+tar
        f(src, name)

    # 一张nii是可以转为一系列dcm图片的
    def niiTOdcm(self, src, tar):
        # 读取图像并转化为array
        itk_image = sitk.ReadImage(src)
        img = sitk.GetArrayFromImage(itk_image)
        # 生成存储路径
        if not os.path.exists(tar):
            os.mkdir(tar)
        print(img.shape)
        # array单帧存储
        for i in range(img.shape[0]):
            out = sitk.GetImageFromArray(img[i, :, :])
            sitk.WriteImage(out, tar+'%04d.dcm'%i)

    # 一系列dcm可以转换为一张nii的
    def dcmTOnii(self, src, tar):
        # 单张间的互转
        if '.dcm' in src:
            itk_image = sitk.ReadImage(src)
            img = sitk.GetArrayFromImage(itk_image)
            out = sitk.GetImageFromArray(img)
            sitk.WriteImage(out, tar)
        # 系列的合并
        else:
            L = Load(src, 'dcm')
            lst = L.load_scan()
            array = []
            for slice_path in lst:
                itk_image = sitk.ReadImage(slice_path)
                img = sitk.GetArrayFromImage(itk_image)
                array.append(img)
            sp = np.shape(array)
            array = np.reshape(array, (sp[0], sp[2], sp[3]))
            print(src, tar, np.shape(array))
            out = sitk.GetImageFromArray(array)
            sitk.WriteImage(out, tar)
# test
# T = Trans()
# T.trans('../dataset/nrrd_heart/nrrd_heart/ID00072637202198161894406_heart.nrrd', 'nii', 'aaaa')
# T.trans('aaaa.nii', 'dcm')